package com.crane.framework.aspect;

import static cn.dev33.satoken.SaManager.log;

import com.crane.framework.annotation.RateLimiter;
import com.crane.framework.enums.LimitType;
import com.crane.framework.exception.RRException;
import com.crane.framework.utils.CommonUtil;
import com.crane.framework.utils.IpUtil;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

/**
 * @author:zhanwang
 * @date: 2023/8/15
 * @time: 18:34
 */
@Aspect
@Component
public class RateLimiterAspect {

  private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);

  @Resource private RedisTemplate<Object, Object> redisTemplate;

  @Resource private RedisScript<Long> redisScript;

  @Before("@annotation(rateLimiter)")
  public void doBefore(JoinPoint point, RateLimiter rateLimiter) {
    String keyPrefix = rateLimiter.keyPrefix();
    int time = rateLimiter.time();
    int count = rateLimiter.count();
    String redisKey = getRedisKey(rateLimiter, point);
    List<Object> keys = Collections.singletonList(redisKey);
    Long number = redisTemplate.execute(redisScript, keys, count, time);
    if (number == null || number.intValue() > count) {
      throw new RRException("访问过于频繁，请稍候再试");
    }
    log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), keyPrefix);
  }

  public String getRedisKey(RateLimiter rateLimiter, JoinPoint point) {
    StringBuilder stringBuffer = new StringBuilder(rateLimiter.keyPrefix());
    if (rateLimiter.limitType() == LimitType.IP) {
      ServletRequestAttributes attributes =
          (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
      if (CommonUtil.isNullOrEmpty(attributes)) {
        return stringBuffer.toString();
      }

      HttpServletRequest request = attributes.getRequest();
      stringBuffer.append(IpUtil.getIpAddr(request)).append("-");
    }
    MethodSignature signature = (MethodSignature) point.getSignature();
    Method method = signature.getMethod();
    Class<?> targetClass = method.getDeclaringClass();
    stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
    return stringBuffer.toString();
  }
}
